DataFrame GroupBy & Aggregation Operations

Summarizing and aggregating grouped datasets in PySpark using aggregate functions like sum, avg, count, min, and max.

What is the GroupBy Operation?

The groupBy() operation segments rows of a DataFrame into unique groups based on values of one or more specified columns. It returns a GroupedData object, which must then be followed by an aggregation operator (such as sum(), avg(), count(), or custom calculations using .agg()) to condense the matching groups back into a standard DataFrame.

Syntax and Multi-Field Aggregations

For single basic aggregates, you can call shorthand methods directly:

# Simple sum of a column grouped by category
df.groupBy("category").sum("sales")

For complex, multi-field calculations, use the .agg() function along with SQL aggregation functions:

from pyspark.sql import functions as F

# Aggregating multiple fields simultaneously with custom column naming
df.groupBy("department").agg(
    F.count("employee_id").alias("employee_count"),
    F.avg("salary").alias("average_salary"),
    F.max("salary").alias("max_salary")
)

Example Usage Pipeline

Below is a complete, copy-paste-ready PySpark script demonstrating advanced groupings and aggregations:

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

# 1. Setup local Spark session
spark = SparkSession.builder \
    .appName("DataFrame GroupBy Demo") \
    .master("local[*]") \
    .getOrCreate()

# 2. Dummy dataset (Sales Transactions)
data = [
    ("Store_A", "Electronics", 1200.0),
    ("Store_A", "Apparel", 250.0),
    ("Store_B", "Electronics", 950.0),
    ("Store_B", "Electronics", 800.0),
    ("Store_A", "Apparel", 150.0),
    ("Store_B", "Apparel", 400.0),
]
columns = ["store_id", "category", "sales_amount"]
df = spark.createDataFrame(data, columns)

# 3. Perform basic sum of sales per store
store_sales = df.groupBy("store_id").sum("sales_amount").withColumnRenamed("sum(sales_amount)", "total_sales")

# 4. Perform complex multi-field aggregation grouped by Store and Category
detailed_agg = df.groupBy("store_id", "category").agg(
    F.count("sales_amount").alias("transaction_count"),
    F.sum("sales_amount").alias("total_sales"),
    F.round(F.avg("sales_amount"), 2).alias("avg_sales")
)

# 5. Show results
print("=== Total Sales Per Store ===")
store_sales.show()

print("=== Store and Category Breakdown ===")
detailed_agg.show(truncate=False)

Rendered Output:

=== Total Sales Per Store ===
+--------+-----------+
|store_id|total_sales|
+--------+-----------+
| Store_A|     1600.0|
| Store_B|     2150.0|
+--------+-----------+

=== Store and Category Breakdown ===
+--------+-----------+-----------------+-----------+---------+
|store_id|category   |transaction_count|total_sales|avg_sales|
+--------+-----------+-----------------+-----------+---------+
|Store_A |Electronics|1                |1200.0     |1200.0   |
|Store_A |Apparel    |2                |400.0      |200.0    |
|Store_B |Electronics|2                |1750.0     |875.0    |
|Store_B |Apparel    |1                |400.0      |400.0    |
+--------+-----------+-----------------+-----------+---------+